using System.Collections;
using System.Collections.Generic;
using UnityEngine;

namespace JufGame
{
	/// <summary>
	/// 神经网络的全连接层
	/// </summary>
    public class Affine : BaseLayer
    {
		public float[] weights, dWeights;
		public float[] bias, dBias;
		public float[] selfInput;
		public ComputeBuffer dWeightBuffer, dBiasBuffer;
		public ComputeBuffer weightBuffer, biasBuffer;
		private ComputeBuffer inputBuffer;
		private int threadGroupsX_Backward, threadGroupsY_Backward;
		private int kernelComputeGrads;
		private int outputCount;

		/// <summary>
		/// 全连接层初始化
		/// </summary>
		/// <param name="shader">用affine compute shader</param>
		/// <param name="batch">输入批次</param>
		/// <param name="inputCount">每批输入的个数</param>
		/// <param name="outputCount">输出个数</param>
		/// <param name="weights">权重矩阵</param>
		/// <param name="bias">偏置矩阵</param>
        public Affine(ComputeShader shader, int batch, int inputCount, int outputCount, float[] weights, float[] bias) : base(shader, batch, inputCount)
        {
			this.outputCount = outputCount;
			this.weights = weights;
			this.bias = bias;
			kernelComputeGrads = compute.FindKernel("ComputeGrads");
			weightBuffer = new ComputeBuffer(weights.Length, sizeof(float));
			biasBuffer = new ComputeBuffer(bias.Length, sizeof(float));
			outputBuffer = new ComputeBuffer(batch * outputCount, sizeof(float));
			threadGroupsX_Backward = Mathf.CeilToInt(outputCount / (THREAD_GROUP_SIZE_X * 1.0f));
			threadGroupsY_Backward = Mathf.CeilToInt(batch / (THREAD_GROUP_SIZE_Y * 1.0f));
        }

        public override void Forward(ComputeBuffer input)
        {
			inputBuffer = input;
			weightBuffer.SetData(weights);
			biasBuffer.SetData(bias);
			compute.SetBuffer(kernelForward, "weightBuffer", weightBuffer);
			compute.SetBuffer(kernelForward, "biasBuffer", biasBuffer);
			compute.SetInt("inputCount", inputCount);
			compute.SetInt("outputCount", outputCount);
			compute.SetInt("batch", batch);
			base.Forward(input);
        }

        public override void Backward(ComputeBuffer dout)
        {
			dBias ??= new float[bias.Length];
			selfInput ??= new float[batch * inputCount];
			dWeights ??= new float[weights.Length];
			dWeightBuffer ??= new ComputeBuffer(dWeights.Length, sizeof(float));
			dBiasBuffer ??= new ComputeBuffer(dBias.Length, sizeof(float));

			biasBuffer.SetData(bias);
			weightBuffer.SetData(weights);

			compute.SetInt("batch", batch);
			compute.SetInt("inputCount", inputCount);
			compute.SetInt("outputCount", outputCount);
			compute.SetBuffer(kernelBackward, "weightBuffer", weightBuffer);

			//反向传播
			dInputBuffer ??= new ComputeBuffer(batch * inputCount, sizeof(float));
			compute.SetBuffer(kernelBackward, "dInputBuffer", dInputBuffer);
			compute.SetBuffer(kernelBackward, "dOutputBuffer", dout);
			compute.Dispatch(kernelBackward, threadGroupsX_Backward, threadGroupsY_Backward, 1);

			//梯度计算
			compute.SetBuffer(kernelComputeGrads, "dOutputBuffer", dout);
			compute.SetBuffer(kernelComputeGrads, "inputBuffer", inputBuffer);
			compute.SetBuffer(kernelComputeGrads, "dWeightBuffer", dWeightBuffer);
			compute.SetBuffer(kernelComputeGrads, "dBiasBuffer", dBiasBuffer);
			compute.Dispatch(kernelComputeGrads, threadGroupsX_Backward, threadGroupsX, 1); //可能有错
			dWeightBuffer.GetData(dWeights);
			dBiasBuffer.GetData(dBias);
		}

        public override void Release()
        {
            base.Release();
			weightBuffer?.Release();
			biasBuffer?.Release();
			dWeightBuffer?.Release();
			dBiasBuffer?.Release();
        }
    }
}
